from pyflink.common import SimpleStringSchema, WatermarkStrategy, Types
from pyflink.datastream import StreamExecutionEnvironment, KeyedProcessFunction, RuntimeContext
from pyflink.datastream.connectors.kafka import KafkaSource, KafkaOffsetsInitializer
from pyflink.datastream.state import ValueStateDescriptor

env = StreamExecutionEnvironment.get_execution_environment()

env.set_parallelism(1)

source = KafkaSource.builder() \
    .set_bootstrap_servers("master:9092") \
    .set_topics("fraud") \
    .set_group_id("my-group") \
    .set_starting_offsets(KafkaOffsetsInitializer.latest()) \
    .set_value_only_deserializer(SimpleStringSchema()) \
    .build()
# 读取数据
lines_ds = env.from_source(source, WatermarkStrategy.no_watermarks(), "Kafka Source")

# 解析数据
frauds_ds = lines_ds.map(lambda line: (line.split(",")[0], float(line.split(",")[1])))

# 对于一个账户，如果出现小于 $1 美元的交易后紧跟着一个大于 $500 的交易，就输出一个报警信息
key_by_ds = frauds_ds.key_by(lambda kv: kv[0])

class FraudKeyedProcessFunction(KeyedProcessFunction):

    def __init__(self):
        self.flag_state = None
        self.min_price_state = None

    def open(self, runtime_context: RuntimeContext):
        self.flag_state = runtime_context.get_state(ValueStateDescriptor("flag", Types.BOOLEAN()))
        self.min_price_state = runtime_context.get_state(ValueStateDescriptor("min_price", Types.FLOAT()))

    def process_element(self, value, ctx: 'KeyedProcessFunction.Context'):
        user_id = value[0]
        price = value[1]

        if self.flag_state.value():
            if price >500:
                print('紧跟着一个大于500的交易')
                min_ptice = self.min_price_state.value()

                yield user_id, min_ptice, price

            self.flag_state.clear()
            self.min_price_state.clear()

        if price < 1:
            print('小于1美元的交易')
            self.flag_state.update(True)
            self.min_price_state.update(price)


key_by_ds.process(FraudKeyedProcessFunction())

env.execute()